#!/bin/sh
# SPDX-License-Identifier: LGPL-2.1+

set -eu
[ -n "${DEBUG:-}" ] && set -x

LXCFSDIR=${LXCFSDIR:-/var/lib/lxcfs}

cg1=$(uuidgen).$$
cg2=$(uuidgen).$$

cleanup() {
	if [ $FAILED -eq 1 ]; then
		exit 1
	fi
	exit 0
}

FAILED=1
trap cleanup EXIT HUP INT TERM

IS_CGROUP_V2=0
grep -qF 'cgroup cgroup' /proc/1/mountinfo || IS_CGROUP_V2=1

if [ "$IS_CGROUP_V2" = "1" ]; then
    [ ! -d /sys/fs/cgroup ] && exit 0
    echo "==> Setting up cgroup"
    mempath=/sys/fs/cgroup/

    memory_limit_file=memory.max
    tasks_file=cgroup.procs
else
    [ ! -d /sys/fs/cgroup/memory ] && exit 0
    echo "==> Setting up memory cgroup"
    initmemory=$(awk -F: '/memory/ { print $3 }' /proc/1/cgroup)
    mempath=/sys/fs/cgroup/memory/${initmemory}

    memory_limit_file=memory.limit_in_bytes
    tasks_file=tasks
fi

rmdir ${mempath}/${cg1} 2>/dev/null || true
rmdir ${mempath}/${cg2} 2>/dev/null || true

echo "==> Testing /proc/meminfo with limit"
mkdir ${mempath}/${cg1}
echo 500000000 > ${mempath}/${cg1}/${memory_limit_file}
echo 1 > ${mempath}/${cg1}/${tasks_file}
m1=$(awk '/^MemTotal:/ { print $2 }' ${LXCFSDIR}/proc/meminfo)

if [ "$IS_CGROUP_V2" = "1" ]; then
    # temporary move to the root cgroup because of
    # "no internal process" constraint
    echo 1 > ${mempath}/${tasks_file}
    echo '+memory' > ${mempath}/${cg1}/cgroup.subtree_control
fi

echo "==> Testing /proc/meminfo with sub-cgroup"
mkdir ${mempath}/${cg1}/${cg2}
echo 1 > ${mempath}/${cg1}/${cg2}/${tasks_file}
m2=$(awk '/^MemTotal:/ { print $2 }' ${LXCFSDIR}/proc/meminfo)

echo "==> Confirming same limits"
[ $m1 -eq $m2 ]

FAILED=0
